import matplotlib.pyplot as plt
import numpy as np
from matplotlib import colormaps
from math_utils import Vector2, Vector2Int
from obstacle import Obstacle
from radiation_field import RadiationField, RadiationSource


class Environment:
    '''路径规划环境类'''

    EMPTY = 0
    '''空白区域标识'''
    OBSTACLE = 1
    '''障碍物标识'''

    def __init__(self, grid_density: int, map: np.ndarray, dose: np.ndarray):
        '''
        初始化环境实例
        Args:
            grid_density: 缩放大小
            map: 环境
            dose: 剂量值
        '''
        self.grid_density = grid_density
        self.map = map
        self.dose = dose
        self.map_shape: Vector2Int = Vector2Int(map.shape[0], map.shape[1])

    @classmethod
    def from_continuous(
        cls,
        obstacle_list: 'list[Obstacle]',
        radiation_field: RadiationField,
        map_size: Vector2Int,
        grid_density: int,
    ) -> 'Environment':
        '''
        从连续环境, 生成环境实例
        Args:
            obstacle_list: 包含环境中所有障碍物的列表
            radiation_field: 环境中的辐射场
            map_size: 地图的尺寸 (宽度, 高度)
            grid_density: 在离散环境中, 每个单位地图长度包含的格子数量数值越大, 地图越精细
        '''
        # 生成网格坐标
        xx, yy = np.meshgrid(
            np.arange(map_size.x * grid_density),
            np.arange(map_size.y * grid_density),
            indexing='ij',
        )
        xx = xx / grid_density
        yy = yy / grid_density

        # 创建地图并标记障碍物
        map = np.zeros([map_size.x * grid_density, map_size.y * grid_density])
        for obstacle in obstacle_list:
            if obstacle.type == obstacle.CIRCLE:
                is_inside = obstacle.point_in_circle(xx, yy)
            elif obstacle.type == obstacle.TRIANGLE:
                is_inside = obstacle.point_in_triangle(xx, yy)
            map = np.maximum(map, is_inside.astype(int))

        # 生成剂量矩阵
        dose = radiation_field.get_dose(xx, yy)

        return cls(grid_density, map, dose)

    def draw(self) -> None:
        '''绘制环境'''
        # 创建画布
        fig = plt.figure(figsize=(6, 6))
        axes = fig.add_subplot(1, 1, 1)
        axes.set_xlim(-0.5, self.map.shape[0] - 0.5)
        axes.set_ylim(-0.5, self.map.shape[1] - 0.5)
        axes.axis('off')

        axes.imshow(self.get_image())
        plt.show()

    def get_image(self) -> np.ndarray:
        '''获取绘制环境的图像数组'''
        # 创建RGB图像
        image = np.zeros((*self.map.shape, 3))

        # 标记障碍物区域
        image[self.map == 1] = (0, 0, 0)

        # 根据剂量值着色空地区域
        cmap = colormaps['plasma']
        mask_0 = (self.map == 0)
        cmap_colors = cmap(self.dose[mask_0])[:, :-1]
        image[mask_0] = cmap_colors

        return image

    def in_map(self, coord: Vector2Int) -> bool:
        '''
        检查坐标是否在地图内
        Args:
            coord: 网格坐标（非实际坐标）
        Returns:
            bool: 在地图内返回 True, 否则返回 False
        '''
        return 0 <= coord.x < self.map.shape[0] and 0 <= coord.y < self.map.shape[1]

    def accessible(self, coord: Vector2Int) -> bool:
        '''
        检查坐标是否可进入
        Args:
            coord: 网格坐标（非实际坐标）
        Returns:
            bool: 可进入返回 True, 否则返回 False
        '''
        return 0 <= coord.x < self.map.shape[0] and 0 <= coord.y < self.map.shape[1] and self.map[coord.x, coord.y] == self.EMPTY


if __name__ == '__main__':
    # 创建测试障碍物
    obstacle_list = [
        Obstacle.circle(Vector2(3, 5), 1),
        Obstacle.triangle(Vector2(2, 2), Vector2(3, 3), Vector2(2, 3)),
        Obstacle.triangle(Vector2(2, 2), Vector2(3, 3), Vector2(3, 2)),
        Obstacle.triangle(Vector2(5, 5), Vector2(8, 8), Vector2(2, 8)),
    ]

    # 创建辐射场
    radiation_field = RadiationField(
        RadiationSource(Vector2(4, 4), intensity=100),
        RadiationSource(Vector2(9, 0), intensity=80),
        RadiationSource(Vector2(0, 9), intensity=60),
    )

    # 测试不同网格密度的环境
    environment_list = [
        Environment.from_continuous(obstacle_list, radiation_field, Vector2Int(10, 10), grid_density=2),
        Environment.from_continuous(obstacle_list, radiation_field, Vector2Int(10, 10), grid_density=4),
        Environment.from_continuous(obstacle_list, radiation_field, Vector2Int(10, 10), grid_density=8),
        Environment.from_continuous(obstacle_list, radiation_field, Vector2Int(10, 10), grid_density=16),
        Environment.from_continuous(obstacle_list, radiation_field, Vector2Int(10, 10), grid_density=32),
    ]

    for environment in environment_list:
        environment.draw()
